-
Notifications
You must be signed in to change notification settings - Fork 14.8k
[SCFToAffine] Add a pass to raise scf to affine ops. #152925
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
@llvm/pr-subscribers-mlir Author: Ming Yan (NexMing) ChangesThis patch supports the conversion from Full diff: https://github.com/llvm/llvm-project/pull/152925.diff 7 Files Affected:
diff --git a/mlir/include/mlir/Conversion/Passes.h b/mlir/include/mlir/Conversion/Passes.h
index 3dc48b2201cf2..2507ef2834dc5 100644
--- a/mlir/include/mlir/Conversion/Passes.h
+++ b/mlir/include/mlir/Conversion/Passes.h
@@ -58,6 +58,7 @@
#include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"
#include "mlir/Conversion/PDLToPDLInterp/PDLToPDLInterp.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
+#include "mlir/Conversion/SCFToAffine/SCFToAffine.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Conversion/SCFToEmitC/SCFToEmitC.h"
#include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 6e1baaf23fcf7..38f35b2dadd94 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1025,6 +1025,18 @@ def ReconcileUnrealizedCastsPass : Pass<"reconcile-unrealized-casts"> {
}];
}
+//===----------------------------------------------------------------------===//
+// SCFToAffine
+//===----------------------------------------------------------------------===//
+
+def RaiseSCFToAffinePass : Pass<"raise-scf-to-affine"> {
+ let summary = "Raise SCF to affine ops";
+ let dependentDialects = [
+ "affine::AffineDialect",
+ "scf::SCFDialect",
+ ];
+}
+
//===----------------------------------------------------------------------===//
// SCFToControlFlow
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h b/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h
new file mode 100644
index 0000000000000..4f87ef8e6c6e4
--- /dev/null
+++ b/mlir/include/mlir/Conversion/SCFToAffine/SCFToAffine.h
@@ -0,0 +1,26 @@
+//===- SCFToAffine.h - SCF to Affine Pass entrypoint ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_
+#define MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_
+
+#include <memory>
+
+namespace mlir {
+class Pass;
+class RewritePatternSet;
+
+#define GEN_PASS_DECL_RAISESCFTOAFFINEPASS
+#include "mlir/Conversion/Passes.h.inc"
+
+/// Collect a set of patterns to convert SCF operations to Affine operations.
+void populateSCFToAffineConversionPatterns(RewritePatternSet &patterns);
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_SCFTOAFFINE_SCFTOAFFINE_H_
diff --git a/mlir/lib/Conversion/CMakeLists.txt b/mlir/lib/Conversion/CMakeLists.txt
index 785cb8293810c..b8059fcbfb028 100644
--- a/mlir/lib/Conversion/CMakeLists.txt
+++ b/mlir/lib/Conversion/CMakeLists.txt
@@ -51,6 +51,7 @@ add_subdirectory(OpenACCToSCF)
add_subdirectory(OpenMPToLLVM)
add_subdirectory(PDLToPDLInterp)
add_subdirectory(ReconcileUnrealizedCasts)
+add_subdirectory(SCFToAffine)
add_subdirectory(SCFToControlFlow)
add_subdirectory(SCFToEmitC)
add_subdirectory(SCFToGPU)
diff --git a/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
new file mode 100644
index 0000000000000..bf1494d6f3cf0
--- /dev/null
+++ b/mlir/lib/Conversion/SCFToAffine/CMakeLists.txt
@@ -0,0 +1,17 @@
+add_mlir_conversion_library(MLIRSCFToAffine
+ SCFToAffine.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/SCFToAffine
+
+ DEPENDS
+ MLIRConversionPassIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRArithDialect
+ MLIRAffineDialect
+ MLIRLLVMDialect
+ MLIRSCFDialect
+ MLIRSCFTransforms
+ MLIRTransforms
+ )
diff --git a/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
new file mode 100644
index 0000000000000..e68bb2123cadc
--- /dev/null
+++ b/mlir/lib/Conversion/SCFToAffine/SCFToAffine.cpp
@@ -0,0 +1,100 @@
+//===- SCFToAffine.cpp - SCF to Affine conversion -------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to raise scf.for, scf.if and loop.terminator
+// ops into affine ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/SCFToAffine/SCFToAffine.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Verifier.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/Passes.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+struct SCFToAffinePass
+ : public impl::RaiseSCFToAffinePassBase<SCFToAffinePass> {
+ void runOnOperation() override;
+};
+
+struct ForOpRewrite : public OpRewritePattern<scf::ForOp> {
+ using OpRewritePattern<scf::ForOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(scf::ForOp op,
+ PatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto lower = op.getLowerBound();
+ auto upper = op.getUpperBound();
+ auto step = op.getStep();
+
+ if (!affine::isValidDim(lower) || !affine::isValidDim(upper) ||
+ !affine::isValidSymbol(step))
+ return llvm::failure();
+
+ auto lowerDim = rewriter.getAffineDimExpr(0);
+ auto upperDim = rewriter.getAffineDimExpr(1);
+ auto stepSym = rewriter.getAffineSymbolExpr(0);
+ auto affineFor = affine::AffineForOp::create(
+ rewriter, loc, ValueRange(), rewriter.getConstantAffineMap(0),
+ ValueRange({lower, upper, step}),
+ AffineMap::get(2, 1,
+ (upperDim - lowerDim + stepSym - 1).floorDiv(stepSym)),
+ 1, op.getInits());
+ auto affineBody = affineFor.getBody();
+
+ if (affineBody->mightHaveTerminator())
+ rewriter.eraseOp(affineBody->getTerminator());
+
+ rewriter.setInsertionPointToStart(affineBody);
+ auto actualIndexMap =
+ AffineMap::get(2, 1, lowerDim + rewriter.getAffineDimExpr(1) * stepSym);
+ Value newIndVar =
+ affine::AffineApplyOp::create(
+ rewriter, op.getLoc(), actualIndexMap,
+ ValueRange({lower, affineFor.getInductionVar(), step}))
+ .getResult();
+
+ SmallVector<Value> argValues;
+ argValues.push_back(newIndVar);
+ llvm::append_range(argValues, affineFor.getRegionIterArgs());
+ rewriter.inlineBlockBefore(op.getBody(), affineBody, affineBody->end(),
+ argValues);
+
+ auto scfYieldOp = cast<scf::YieldOp>(affineBody->getTerminator());
+ rewriter.setInsertionPointToEnd(affineBody);
+ rewriter.replaceOpWithNewOp<affine::AffineYieldOp>(
+ scfYieldOp, scfYieldOp->getOperands());
+
+ rewriter.replaceOp(op, affineFor);
+ return success();
+ }
+};
+
+} // namespace
+
+void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) {
+ patterns.add<ForOpRewrite>(patterns.getContext());
+}
+
+void SCFToAffinePass::runOnOperation() {
+ RewritePatternSet patterns(&getContext());
+ populateSCFToAffineConversionPatterns(patterns);
+
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+ signalPassFailure();
+}
diff --git a/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
new file mode 100644
index 0000000000000..6f419bc8ee9ce
--- /dev/null
+++ b/mlir/test/Conversion/SCFToAffine/scf-to-affine.mlir
@@ -0,0 +1,34 @@
+// RUN: mlir-opt -raise-scf-to-affine -split-input-file %s | FileCheck %s
+
+// CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0, s1, s2] -> ((s0 - s1 + s2 - 1) floordiv s0)>
+// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
+// CHECK-LABEL: func.func @simple_loop(
+// CHECK-SAME: %[[ARG0:.*]]: memref<?xi32>,
+// CHECK-SAME: %[[ARG1:.*]]: memref<3xindex>) {
+// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i32
+// CHECK: %[[VAL_1:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 2 : index
+// CHECK: %[[VAL_4:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_1]]] : memref<3xindex>
+// CHECK: %[[VAL_5:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_2]]] : memref<3xindex>
+// CHECK: %[[VAL_6:.*]] = memref.load %[[ARG1]]{{\[}}%[[VAL_3]]] : memref<3xindex>
+// CHECK: affine.for %[[VAL_7:.*]] = 0 to #[[$ATTR_0]](){{\[}}%[[VAL_6]], %[[VAL_4]], %[[VAL_5]]] {
+// CHECK: %[[VAL_8:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_4]], %[[VAL_7]]){{\[}}%[[VAL_6]]]
+// CHECK: memref.store %[[VAL_0]], %[[ARG0]]{{\[}}%[[VAL_8]]] : memref<?xi32>
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+func.func @simple_loop(%arg0: memref<?xi32>, %arg1: memref<3xindex>) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %0 = memref.load %arg1[%c0] : memref<3xindex>
+ %1 = memref.load %arg1[%c1] : memref<3xindex>
+ %2 = memref.load %arg1[%c2] : memref<3xindex>
+ scf.for %arg2 = %0 to %1 step %2 {
+ memref.store %c0_i32, %arg0[%arg2] : memref<?xi32>
+ }
+ return
+}
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please update the commit message to explain why this is added, https://mlir.llvm.org/getting_started/Contributing/#commit-messages. It is also a good idea to describe whether this a one-off patch or whether further work is planned.
230fa93
to
63039ce
Compare
This patch supports the conversion from `scf.for` to `affine.for`.
63039ce
to
48dad4a
Compare
This patch adds an initial pass to convert eligible
scf.for
loops intoaffine.for
. The Affine dialect enables stronger optimization and analysis capabilities compared to SCF, including dependence analysis, loop tiling/fusion, and parallelization. Promoting analyzable SCF loops to Affine provides opportunities for polyhedral-style optimizations at the core dialect level.This SCF-to-Affine pass is inspired by the Polygeist project, which implemented similar transformations. However, this implementation is independently developed in MLIR project. This allows frontends generating SCF loops, such as ClangIR or Flang, to benefit from Affine-based transformations.
This patch is the first step in an incremental effort: